package 红黑树;

import sun.net.idn.Punycode;

/**
 * @author:谢君臣
 * @Date:2021/8/26-9:22
 * @version:1.0
 * @function:
 */
public class RBTree {

    private TreeNode root;
    private int size;

    class TreeNode {
        public Color color;
        public int val;
        public TreeNode parent;
        public TreeNode left;
        public TreeNode right;

        public TreeNode() {
        }

        public TreeNode(int val) {
            this.color = Color.R;
            this.val = val;
        }
    }

    enum Color {
        R, B;
    }

    public RBTree() {
    }

    public void insert(int val) {
        TreeNode node = new TreeNode(val);
        if (root == null) {  // 添加根节点
            node.color = Color.B;
            root = node;
        } else {
            //找到待插入位置
            TreeNode treeNode = findIndex(val);
            if (treeNode == null) return; // 找到了已存在的节点
            if (treeNode.val > val) treeNode.left = node;
            else treeNode.right = node;
            node.parent = treeNode;
            BalanceInsert(node);
        }
        size++;
    }

    private void BalanceInsert(TreeNode node) {
        if (node.parent.color == Color.B) // 如果父节点为黑色不处理，如果为红色需要调整平衡
            return;
        TreeNode parent = node.parent;
        TreeNode grand = parent.parent;
        if (grand.right == null || grand.right.color == Color.B||grand.left == null || grand.left.color == Color.B) { // 叔伯节点为黑色
            if (grand.left == parent && node == parent.right) // < 型的结构
                leftRotate(parent);
            if (grand.left == parent && node == parent.left) // / 型的结构
                ;
            if (grand.right == parent && node == parent.left) // > 型的结构
                ;
            if (grand.right == parent && node == parent.right) // \ 型的结构
                ;
        } else {  // 叔伯节点为红色
            grand.color = Color.R;
            grand.right.color = Color.B;
            parent.color = Color.B;
            if (grand == root) {
                grand.color = Color.B;  // 树的高度+1
            } else
                BalanceInsert(grand);  // 递归向上
        }
    }

    private void leftRotate(TreeNode node) {
        TreeNode parent = node.parent;
        if (parent==null){
            root = node.right;root.parent=null;
            node.parent = root;
            if (root.left!=null) {
                node.right = root.left;
                root.left = node;
                node.right.parent = node;
            }
        }else {
            node.parent.left = node.right;
            node.right = node.parent.left.left;
            node.parent.left.parent = node.parent;
            node.parent = node.parent.left;
            node.parent.left = node;
            if (node.right!=null)
                node.right.parent = node;
        }
    }

    public TreeNode findIndex(int val) {
        TreeNode temp = root, parent = null;
        while (temp != null) {
            parent = temp;
            if (parent.val > val) temp = parent.left;
            else if (parent.val < val) temp = parent.right;
            else return null;
        }
        return parent;
    }

    public TreeNode getRoot() {
        return root;
    }

    public int getSize() {
        return size;
    }
}
